// 到这个网址可以在线编码，实操transformer.js，并且相关模型已经准备好：https://chn.ai/embedding.html
// 视频讲解 https://www.bilibili.com/video/BV1ZD421j7os/
// 文字讲解：https://blog.csdn.net/fribbler/article/details/136728513


import { AutoModel, AutoTokenizer, dot } from 'https://res.chn.ai/module/transformers@2.15.1/transformers@2.15.1.js'

let tokenizer = await AutoTokenizer.from_pretrained('Xenova/jina-embeddings-v2-base-zh')
let model = await AutoModel.from_pretrained('Xenova/jina-embeddings-v2-base-zh')

function stringify(obj) {
    return JSON.stringify(obj, (key, value) => {
        return typeof value === 'bigint'
            ? value.toString()
            : value
        
    }, null, 2)
}


var sentences = ['苹果这次发布会水分很大', '一个苹果里面含有丰富的水分', '保持土壤水分', '水分', '苹果']
var tokens = await tokenizer(sentences, {padding: true})
var tokenIds = Array.from(tokens.input_ids.data)
var words = await tokenizer.decode(tokenIds).split(' ')

console.log(words.map((w, i) => {
	return `${w}(${tokenIds[i]})`
}).join(' '))
// <s>(0) 苹果(27691) 这次(27937) 发布会(40770) 水分(37459) 很大(29075) </s>(2) <pad>(1) <s>(0) 一个(21752) 苹果(27691) 里面(29468) 含有(31173) 丰富的(28195) 水分(37459) </s>(2)

var embeddings = (await model(tokens)).last_hidden_state
var vectorLen = embeddings.dims[2]
var embeddingData = Array.from(embeddings.data)
var wordEmbeddings = []
for(var i = 0; i < embeddingData.length; i += vectorLen) {
	var wordEmbedding = embeddingData.slice(i, i + vectorLen)
	// console.log(`${words[i/vectorLen]} ${wordEmbedding.length}`)
	wordEmbeddings.push({index: i/vectorLen, word: words[i/vectorLen], wordEmbedding})
}
console.log(stringify(wordEmbeddings))

//苹果
var tt = {1: '苹果这次发布会水分很大', 10: '一个苹果里面含有丰富的水分', 33: '苹果'}
for(var i1 of [1, 10, 33]) {
	for(var i2 of [1, 10, 33]) {
		console.log(`${tt[i1]} - ${tt[i2]} ${Math.floor(dot(wordEmbeddings[i1].wordEmbedding, wordEmbeddings[i2].wordEmbedding) * 10) / 10}`)
	}
}

